'''
自适应加权SG滤波
'''
import os
import time

import numpy as np
from osgeo import gdal, osr
from osgeo.gdalconst import GRIORA_NearestNeighbour, GDT_UInt16
from scipy.interpolate import interp1d
from scipy.signal import savgol_filter

np.seterr(divide='ignore', invalid='ignore')  # 消除被除数为0的警告


# 读取tif文件
def read_tif(file_path):
    dataset = gdal.Open(file_path)
    data = dataset.ReadAsArray()
    geotransform = dataset.GetGeoTransform()
    return data, geotransform


# 写入tif文件
def write_tif(file_path, data, geotransform, projection, nodata, dataType):
    gdal_type = ''
    if dataType == 'int16':
        gdal_type = gdal.GDT_Int16
    elif dataType == 'int32':
        gdal_type = gdal.GDT_Int32
    elif dataType == 'float32':
        gdal_type = gdal.GDT_Float32
    elif dataType == 'float64':
        gdal_type = gdal.GDT_Float64
    elif dataType == 'byte':
        gdal_type = gdal.GDT_Byte
    elif dataType == 'uint16':
        gdal_type = gdal.GDT_UInt16
    elif dataType == 'uint32':
        gdal_type = gdal.GDT_UInt32
    driver = gdal.GetDriverByName("GTiff")
    rows, cols = data.shape
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    dataset.SetGeoTransform(geotransform)
    # 定义投影
    # prj = osr.SpatialReference()
    # prj.ImportFromEPSG(4326)
    dataset.SetProjection(projection)
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    del dataset
    print("写入成功：{}".format(file_path))


def savitzky_golay_filter(signal, window_length, polyorder, weights=None):
    if window_length % 2 != 1:
        raise ValueError("窗口长度必须为奇数")
    if window_length < polyorder + 2:
        raise ValueError("窗口长度必须大于等于多项式阶数加2")
    half_window = window_length // 2

    if weights is None:
        weights = np.ones(signal.shape)

    # 将signal后半个窗口补充到前面，前半个窗口补充到后面
    signal = np.concatenate((signal[-half_window:], signal, signal[:half_window]))
    weights = np.concatenate((weights[-half_window:], weights, weights[:half_window]))
    # print(signal)
    # print(weights)
    filtered_signal = np.zeros_like(signal)

    for i in range(half_window, len(signal) - half_window):
        if weights[i] == 1:
            filtered_signal[i] = signal[i]
        else:
            segment = signal[i - half_window: i + half_window + 1]
            weight = weights[i - half_window: i + half_window + 1]

            filtered = np.where(np.isnan(np.dot(weight, segment) / np.sum(weight)), -3000, np.dot(weight, segment) / np.sum(weight))
            filtered_signal[i] = np.where(np.isinf(filtered), -3000, filtered)

        return filtered_signal[half_window:len(signal) - half_window]


def IAW_SG_HDF(HDF_dir, Output_dir):
    start = time.time()
    # 读取文件夹中所有tif后缀的文件路径存成一个列表
    HDF_path_list = [HDF_dir + '\\' + file for file in os.listdir(HDF_dir) if file.endswith('.hdf')]
    print(HDF_path_list[0])
    print(HDF_path_list[-1])
    # 根据文件数量创建一个时间数组，步长为16
    start_time = int(HDF_path_list[0].split('\\')[-1].split('.')[1][-3:])
    end_time = int(HDF_path_list[-1].split('\\')[-1].split('.')[1][-3:])
    # print(start_time, end_time)
    # 根据开始和结束时间创建一个时间数组，步长为16，包括结束时间
    VI_time_array = np.arange(start_time, end_time + 1, 16)
    print(VI_time_array)
    # 插值时间数组，步长为1
    VI_interp_array = np.arange(start_time, end_time + 1, 1)
    # print(VI_interp_array)

    # 读取单幅影像获取一个时间点的图像信息
    dataset = gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_16DAY_1km_VI:"1 km 16 days EVI"'.format(HDF_path_list[0]))
    rows, cols = dataset.RasterYSize, dataset.RasterXSize
    # print(rows, cols)
    geotransform = dataset.GetGeoTransform()
    projection = dataset.GetProjection()
    nodata = dataset.GetRasterBand(1).GetNoDataValue()
    # print(nodata)
    # 设置一个空数组存放所有图像数据
    VI_array = np.zeros((len(HDF_path_list), rows, cols)).astype(np.int16)
    Pixel_Reliability_array = np.zeros((len(HDF_path_list), rows, cols)).astype(np.int8)
    Quality_array = np.zeros((len(HDF_path_list), rows, cols)).astype(np.uint16)
    # 循环读取所有图像数据
    for i in range(len(HDF_path_list)):
        dataset = gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_16DAY_1km_VI:"1 km 16 days EVI"'.format(HDF_path_list[i]))
        VI_array[i, :, :] = dataset.ReadAsArray()
        dataset = gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_16DAY_1km_VI:"1 km 16 days pixel reliability"'.format(HDF_path_list[i]))
        Pixel_Reliability_array[i, :, :] = dataset.ReadAsArray(buf_xsize=cols, buf_ysize=rows, resample_alg=GRIORA_NearestNeighbour)
        dataset = gdal.Open(r'HDF4_EOS:EOS_GRID:"{}":MODIS_Grid_16DAY_1km_VI:"1 km 16 days VI Quality"'.format(HDF_path_list[i]))
        # 读取数组并设置类型为16位无符号整数
        Quality_array[i, :, :] = dataset.ReadAsArray(buf_xsize=cols, buf_ysize=rows, buf_type=GDT_UInt16, resample_alg=GRIORA_NearestNeighbour)
        # 将质量控制数据中数据转为2进制
        binary_arr = np.vectorize(np.binary_repr)(Quality_array[i, :, :])
        # 将上面二维数组每个元素字符串后6位转为十进制组成新的数组
        decimal_arr = np.array([int(binary[-6:], 2) for binary in binary_arr.flatten()]).reshape(binary_arr.shape)
        Quality_array[i, :, :] = decimal_arr.astype(np.uint8)
        dataset = None
    print(VI_array.shape)
    # 切换维度，便于后面提取每个像素对应各个时期值
    VI_array = np.transpose(VI_array, (1, 2, 0))
    Pixel_Reliability_array = np.transpose(Pixel_Reliability_array, (1, 2, 0))
    Quality_array = np.transpose(Quality_array, (1, 2, 0))
    print(VI_array.shape)
    # 建立权重数组
    Weight_array = np.zeros(VI_array.shape)
    for i in range(VI_array.shape[0]):
        for j in range(VI_array.shape[1]):
            weight_temp = np.zeros(VI_array[i, j, :].shape)
            weight_temp[Pixel_Reliability_array[i, j, :] == 0] = 1
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 0)] = 1
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 1)] = 0.8
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 2)] = 0.6
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 3)] = 0.4
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 4)] = 0.2
            Weight_array[i, j, :] = weight_temp
    # SG滤波结果保存
    SG_filter_arr = np.zeros((VI_array.shape[0], VI_array.shape[1], VI_interp_array.shape[0]))
    for i in range(VI_array.shape[0]):
        for j in range(VI_array.shape[1]):
            # 第一次滤波
            SG_filter_one = savitzky_golay_filter(VI_array[i, j, :], window_length=7, polyorder=2, weights=Weight_array[i, j, :])
            # 更新权重，第一次滤波结果和原始数据对比，小于滤波结果的权重赋1，大于滤波结果的赋两者比值，要考虑正负，权重都在0-1之间
            weight_new = SG_filter_one / VI_array[i, j, :]
            weight_new[np.abs(SG_filter_one) >= np.abs(VI_array[i, j, :])] = 1
            # 初始权重为0的也设0
            weight_new[Weight_array[i, j, :] == 0] = 0
            SG_filter_second = savitzky_golay_filter(SG_filter_one, window_length=7, polyorder=2, weights=weight_new)
            # 进行内插
            interp_func = interp1d(VI_time_array, SG_filter_second, kind='quadratic')
            SG_filter_arr[i, j, :] = interp_func(VI_interp_array)
    # 维度转换回去
    SG_filter_arr = np.transpose(SG_filter_arr, (2, 0, 1))
    print(SG_filter_arr.shape)
    # 保存
    for i in range(SG_filter_arr.shape[0]):
        # 按照前面的列表设置输出文件名
        Output_file_name = Output_dir + '\\' + str(VI_interp_array[i]) + '_SG_filter.tif'
        write_tif(Output_file_name, SG_filter_arr[i, :, :], geotransform, projection, nodata)
    end = time.time()
    print('SG filter time:', end - start)
    print('SG filter finished!')


def IAW_SG_TIF(EVI_dir, PR_dir, QC_dir, Output_dir):
    start = time.time()
    # 读取文件夹中所有tif后缀的文件路径存成一个列表
    EVI_path_list = [EVI_dir + '\\' + file for file in os.listdir(EVI_dir) if file.endswith('.tif')]
    PR_path_list = [PR_dir + '\\' + file for file in os.listdir(PR_dir) if file.endswith('.tif')]
    QC_path_list = [QC_dir + '\\' + file for file in os.listdir(QC_dir) if file.endswith('.tif')]
    # 根据文件数量创建一个时间数组，步长为16
    start_time = int(EVI_path_list[0].split('\\')[-1].split('.')[0][-3:])
    end_time = int(EVI_path_list[-1].split('\\')[-1].split('.')[0][-3:])
    # print(start_time, end_time)
    # 根据开始和结束时间创建一个时间数组，步长为16，包括结束时间
    VI_time_array = np.arange(start_time, end_time + 1, 16)
    print(VI_time_array)
    # 插值时间数组，步长为1
    VI_interp_array = np.arange(start_time, end_time + 1, 1)
    # print(VI_interp_array)

    # 读取单幅影像获取一个时间点的图像信息
    dataset = gdal.Open(EVI_path_list[0])
    rows, cols = dataset.RasterYSize, dataset.RasterXSize
    # print(rows, cols)
    geotransform = dataset.GetGeoTransform()
    projection = dataset.GetProjection()
    nodata = dataset.GetRasterBand(1).GetNoDataValue()
    # print(nodata)
    # 设置一个空数组存放所有图像数据
    VI_array = np.zeros((len(EVI_path_list), rows, cols)).astype(np.int16)
    Pixel_Reliability_array = np.zeros((len(EVI_path_list), rows, cols)).astype(np.int8)
    Quality_array = np.zeros((len(EVI_path_list), rows, cols)).astype(np.uint16)
    # 循环读取所有图像数据
    for i in range(len(EVI_path_list)):
        dataset = gdal.Open(EVI_path_list[i])
        VI_array[i, :, :] = dataset.ReadAsArray()
        dataset = gdal.Open(PR_path_list[i])
        Pixel_Reliability_array[i, :, :] = dataset.ReadAsArray()
        dataset = gdal.Open(QC_path_list[i])
        # 读取数组并设置类型为16位无符号整数
        Quality_array[i, :, :] = dataset.ReadAsArray()
        # 将质量控制数据中数据转为2进制
        binary_arr = np.vectorize(np.binary_repr)(Quality_array[i, :, :])
        # 将上面二维数组每个元素字符串后6位转为十进制组成新的数组
        decimal_arr = np.array([int(binary[-6:], 2) for binary in binary_arr.flatten()]).reshape(binary_arr.shape)
        Quality_array[i, :, :] = decimal_arr.astype(np.uint8)
        dataset = None
    print(VI_array.shape)
    # 切换维度，便于后面提取每个像素对应各个时期值
    VI_array = np.transpose(VI_array, (1, 2, 0))
    Pixel_Reliability_array = np.transpose(Pixel_Reliability_array, (1, 2, 0))
    Quality_array = np.transpose(Quality_array, (1, 2, 0))
    print(VI_array.shape)
    # 建立权重数组
    Weight_array = np.zeros(VI_array.shape)
    for i in range(VI_array.shape[0]):
        for j in range(VI_array.shape[1]):
            weight_temp = np.zeros(VI_array[i, j, :].shape)
            weight_temp[Pixel_Reliability_array[i, j, :] == 0] = 1
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 0)] = 1
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 1)] = 0.8
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 2)] = 0.6
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 3)] = 0.4
            weight_temp[(Pixel_Reliability_array[i, j, :] == 1) + (Quality_array[i, j, :] == 4)] = 0.2
            Weight_array[i, j, :] = weight_temp
    # SG滤波结果保存
    SG_filter_arr = np.zeros((VI_array.shape[0], VI_array.shape[1], VI_interp_array.shape[0]))
    for i in range(VI_array.shape[0]):
        for j in range(VI_array.shape[1]):
            # 第一次滤波
            SG_filter_one = savitzky_golay_filter(VI_array[i, j, :], window_length=7, polyorder=2, weights=Weight_array[i, j, :])
            # 更新权重，第一次滤波结果和原始数据对比，小于滤波结果的权重赋1，大于滤波结果的赋两者比值，要考虑正负，权重都在0-1之间
            weight_new = SG_filter_one / VI_array[i, j, :]
            weight_new[np.abs(SG_filter_one) >= np.abs(VI_array[i, j, :])] = 1
            # 初始权重为0的也设0
            weight_new[Weight_array[i, j, :] == 0] = 0
            SG_filter_second = savitzky_golay_filter(SG_filter_one, window_length=7, polyorder=2, weights=weight_new)
            # 进行内插
            interp_func = interp1d(VI_time_array, SG_filter_second, kind='quadratic')
            SG_filter_arr[i, j, :] = interp_func(VI_interp_array)
    # 维度转换回去
    SG_filter_arr = np.transpose(SG_filter_arr, (2, 0, 1))
    print(SG_filter_arr.shape)
    # 保存
    for i in range(SG_filter_arr.shape[0]):
        # 按照前面的列表设置输出文件名
        Output_file_name = Output_dir + '\\' + str(VI_interp_array[i]) + '_SG_filter.tif'
        write_tif(Output_file_name, SG_filter_arr[i, :, :], geotransform, projection, nodata, dataType='uint16')
    end = time.time()
    print('SG filter time:', end - start)
    print('SG filter finished!')


if __name__ == '__main__':
    # HDF_dir = r'G:\MOD13A2\h25v05'
    # Output_dir = r'G:\MOD13A2\OUT_h25v05'
    # HDF_dir = r'G:\MOD13A2\h26v05'
    # Output_dir = r'G:\MOD13A2\OUT_h26v05'
    # IAW_SG_HDF(HDF_dir, Output_dir)

    EVI_dir = r'G:\test\process_result\EVI_tif\merge_clip\2013\EVI'
    PR_dir = r'G:\test\process_result\EVI_tif\merge_clip\2013\PR'
    QC_dir = r'G:\test\process_result\EVI_tif\merge_clip\2013\QC'
    Output_dir = r'G:\test\process_result\EVI_tif\EVI_SG_2013'
    IAW_SG_TIF(EVI_dir, PR_dir, QC_dir, Output_dir)
